import torch
import torchvision.transforms as T
from torch.utils.data import DataLoader

from timm.data.random_erasing import RandomErasing

from .dataset import UcfFramesDataSet

__factory = {
    'ucf101': UcfFramesDataSet,
}


def train_collate_fn(batch):
    """
    # collate_fn这个函数的输入就是一个list，list的长度是一个batch size，list中的每个元素都是__getitem__得到的结果
    """
    imgs, pids, camids, viewids, _ = zip(*batch)
    pids = torch.tensor(pids, dtype=torch.int64)
    viewids = torch.tensor(viewids, dtype=torch.int64)
    camids = torch.tensor(camids, dtype=torch.int64)
    return torch.stack(imgs, dim=0), pids, camids, viewids,


def val_collate_fn(batch):
    imgs, pids, camids, viewids, img_paths = zip(*batch)
    viewids = torch.tensor(viewids, dtype=torch.int64)
    camids_batch = torch.tensor(camids, dtype=torch.int64)
    return torch.stack(imgs, dim=0), pids, camids, camids_batch, viewids, img_paths


def make_dataloader(cfg):
    train_transforms = T.Compose([
        T.Resize(cfg.INPUT.SIZE_TRAIN, interpolation=3),
        T.RandomHorizontalFlip(p=cfg.INPUT.PROB),
        T.Pad(cfg.INPUT.PADDING),
        T.RandomCrop(cfg.INPUT.SIZE_TRAIN),
        T.ToTensor(),
        T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD),
        RandomErasing(probability=cfg.INPUT.RE_PROB, mode='pixel', max_count=1, device='cpu'),
    ])

    val_transforms = T.Compose([
        T.Resize(cfg.INPUT.SIZE_TEST),
        T.ToTensor(),
        T.Normalize(mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD)
    ])

    num_workers = cfg.DATALOADER.NUM_WORKERS
    batch_size = cfg.SOLVER.BATCH_SIZE
    pin_memory = True

    train_dataset = __factory[cfg.DATASETS.NAMES](ann_file=cfg.DATASETS.ANN_FILE_TRAIN,
                                                  root_dir=cfg.DATASETS.ROOT_DIR,
                                                  t_num=cfg.SOLVER.IMS_PER_BATCH,
                                                  transform=train_transforms,
                                                  )

    val_dataset = __factory[cfg.DATASETS.NAMES](ann_file=cfg.DATASETS.ANN_FILE_VAL,
                                                root_dir=cfg.DATASETS.ROOT_DIR,
                                                t_num=cfg.TEST.IMS_PER_BATCH,
                                                transform=val_transforms,
                                                )

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        num_workers=num_workers,
        batch_size=batch_size,
        pin_memory=True,
    )

    val_loader = DataLoader(
        val_dataset, batch_size=cfg.TEST.TEST_BATCH, shuffle=False, num_workers=num_workers,
        pin_memory=pin_memory, drop_last=False,
        # collate_fn=val_collate_fn
    )

    return train_loader, val_loader
